import json
import os

def process_tags(input_file, output_file, removed_file):
    """
    处理JSON文件中的标签：
    - 若tag仅包含"幼儿园卫生间"，则保留
    - 若tag包含"幼儿园卫生间"且有其他标签，则移除该标签
    - 保存处理后的数据和被修改的记录
    """
    # 读取原始JSON数据
    if not os.path.exists(input_file):
        raise FileNotFoundError(f"输入文件不存在: {input_file}")
    
    with open(input_file, 'r', encoding='utf-8') as f:
        data = json.load(f)  # 假设输入为标准JSON数组格式

    processed_data = []
    removed_entries = []
    
    for item in data:
        # 跳过缺少必要字段的条目
        if "tag" not in item or "image" not in item:
            processed_data.append(item)
            continue
            
        tag_str = item["tag"].strip()
        if "幼儿园卫生间" not in tag_str:
            # 不包含目标标签，直接保留
            processed_data.append(item)
            continue
        
        # 分割标签并去重
        tags = [t.strip() for t in tag_str.split(',') if t.strip()]
        tag_set = list(set(tags))  # 去重
        
        if len(tag_set) == 1 and tag_set[0] == "幼儿园卫生间":
            # 仅含目标标签，保留
            processed_data.append(item)
        else:
            # 移除目标标签
            new_tags = [t for t in tag_set if t != "幼儿园卫生间"]
            # 记录被修改的条目
            removed_entries.append({
                "image": item["image"],
                "original_tag": tag_str,
                "new_tag": ",".join(new_tags)
            })
            # 更新当前条目并保存
            item["tag"] = ",".join(new_tags)
            processed_data.append(item)
    
    # 保存处理结果
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(processed_data, f, ensure_ascii=False, indent=2)
    
    # 保存被移除的记录
    with open(removed_file, 'w', encoding='utf-8') as f:
        json.dump(removed_entries, f, ensure_ascii=False, indent=2)
    
    print(f"处理完成：")
    print(f"原始数据量: {len(data)}")
    print(f"修改条目数: {len(removed_entries)}")
    print(f"处理后文件: {output_file}")
    print(f"修改记录文件: {removed_file}")

if __name__ == "__main__":
    # 请根据实际文件路径修改
    INPUT_JSON = "processed_NewGZ_tag/train/cleaned_train.json"       # 原始JSON文件路径
    OUTPUT_JSON = "processed_NewGZ_tag/train/cleaned_train2.json"     # 处理后的数据
    REMOVED_JSON = "processed_NewGZ_tag/train/removed.json"   # 被修改的记录
    
    process_tags(INPUT_JSON, OUTPUT_JSON, REMOVED_JSON)
    